# 导入相关模块
import torch
from torch.utils.data import TensorDataset
# 导入构造数据集对象
from torch.utils.data import DataLoader
# 导入数据加载器
from torch import nn
# nn模块中有平方损失函数和假设函数
from torch import optim
# optim模块中有优化器
from sklearn.datasets import make_regression
# 创建线性回归模型数据集
import matplotlib.pyplot as plt

# 绘图模块
#  pip install scikit-learn matplotlib -i https://pypi.tuna.tsinghua.edu.cn/simple


# 使用PyTorch的nn.SELoss0代替自定义的平方损失函数
# 使用PyTorch的data.DataLoader代替自定义的数据加载器
# 使用PyTorch的optim.SGD代替自定义的优化器
# 使用PyTorch的nn.Linear代替自定义的假设函数


plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def create_dataset():
    x, y, conf = make_regression(n_samples=100,
                                 n_features=1,
                                 random_state=0,
                                 noise=10,
                                 bias=10,
                                 coef=True)

    # 将构建数据转换为张量类型
    x = torch.tensor(x, dtype=torch.float32)
    y = torch.tensor(y, dtype=torch.float32)
    conf = torch.tensor(conf, dtype=torch.float32)
    return x, y, conf


if __name__ == '__main__':
    # 生成的数据
    x, y, conf = create_dataset()
    # 绘制数据的真实线性回归结果
    # plt.scatter(x, y)
    # x_new = torch.linspace(x.min(), x.max(), 1000)
    # y1 = x_new * conf + 1.5
    # # 使用新生成的 x_new 和 y1 绘图
    # plt.plot(x_new, y1, label='real')
    # plt.grid()
    # plt.legend()
    # # 切换 Matplotlib 后端， 使用TkAgg交互式窗口避免代码运行完后窗口直接被关闭
    # plt.switch_backend('TkAgg')
    # plt.show()

    # 构造数据集对象
    dataset = TensorDataset(x, y)
    # 构造数据加载器
    # dataset: 数据集对象
    # batch_size: 批量训练样本数据
    # shuffle: 样本数据是否进行乱序
    data_loader = DataLoader(dataset, batch_size=10, shuffle=True)
    # 构造线性回归模型
    # in_features: 指的是输入张量的大小size
    # out_features: 指的是输出张量的大小size
    model = nn.Linear(in_features=1, out_features=1)
    # 构造平方损失函数
    loss = nn.MSELoss()
    # 构造优化器
    optimizer = optim.SGD(params=model.parameters(), lr=0.001)




    epochs = 100
    # 损失的变化
    loss_epoch = []
    total_loss = 0.0
    train_sample = 0.0
    for _ in range(epochs):
        for train_x, train_y in data_loader:
            # 将一个batch的训练数据输入模型
            y_pred = model(train_x.type(torch.float32))
            # 计算损失值
            loss_values = loss(y_pred, train_y.reshape(-1, 1).type(torch.float32))
            total_loss += loss_values
            train_sample += len(train_y)  # 计算训练样本的数量
            optimizer.zero_grad()   # 梯度清零
            loss_values.backward()  # 自动微分（反向传播）
            optimizer.step()  # 优化器更新参数
            loss_epoch.append(total_loss / train_sample)  # 获取每个batch的损失


    # 绘制损失值变化曲线
    plt.plot(range(epochs), loss_epoch)
    plt.show()

    plt.scatter(x, y)
    x1 = torch.linspace(x.min(), x.max(), 1000)
    y0 = torch.tensor([v*model.weight+model.bias for v in x1])
    y1 = torch.tensor([v*conf+1.5 for v in x1])
    plt.plot(x1, y0, label='yuce')
    plt.plot(x1, y1, label='zhenshi')
    plt.grid()
    plt.legend()
    plt.show()

